from transformers import MarianMTModel, MarianTokenizer
import os
os.environ['HTTPS_PROXY'] = 'https://127.0.0.1:7890'
os.environ['HTTP_PROXY'] = 'http://127.0.0.1:7890'

# '(MaxRetryError("HTTPSConnectionPool(host='huggingface.co', port=443): Max retries exceeded with url: /Helsinki-NLP/opus-mt-en-zh/resolve/main/tokenizer_config.json (Caused by ProxyError('Unable to connect to proxy', SSLError(SSLEOFError(8, 'EOF occurred in violation of protocol (_ssl.c:997)'))))"), '(Request ID: 730649d5-6f71-428e-a23c-03788aaace97)')' thrown while requesting HEAD https://huggingface.co/Helsinki-NLP/opus-mt-en-zh/resolve/main/tokenizer_config.json
# Retrying in 1s [Retry 1/5].
class LocalTranslator:
    def __init__(self, source_lang="en", target_lang="zh"):
        self.model_name = f"Helsinki-NLP/opus-mt-{source_lang}-{target_lang}"
        self.tokenizer = MarianTokenizer.from_pretrained(self.model_name)
        self.model = MarianMTModel.from_pretrained(self.model_name)

    def translate(self, text):
        # 编码输入文本
        inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)

        # 生成翻译
        translated = self.model.generate(**inputs)

        # 解码输出
        translated_text = self.tokenizer.decode(translated[0], skip_special_tokens=True)
        return translated_text


# 使用示例
translator = LocalTranslator("en", "zh")
result = translator.translate("This is a test sentence.")
print(result)